package main

import (
	"fmt"
	"math/rand"
	"time"
)

func main() {
	arr := []int{0, 0, 1, 2, 4, 2, 2, 3, 1, 4}
	a := getLeastNumbers(arr, 8)
	fmt.Println(a)
}

func getLeastNumbers(arr []int, k int) []int {
	if len(arr) < k || k <= 0 {
		return []int{}
	}
	location := 0
	start, end := 0, len(arr)-1
	for location != k-1 {
		location = partition(arr, start, end)
		if location < k-1 {
			start = location + 1
		}
		if location > k-1 {
			end = location - 1
		}
	
	}
	ret := []int{}
	for i := 0; i <= location; i++ {
		ret = append(ret, arr[i])
	}
	return ret
}

func partition(arr []int, start, end int) int {
	if len(arr) == 0 || start >= len(arr) || start < 0 {
	// 	panic("invalid input")
	// }
	if start >= end {
		return start
	}

	rand := rand.Intn(end-start) + start
	arr[rand], arr[end] = arr[end], arr[rand]
	small := start

	for i := start; i < end; i++ {
		if arr[i] >= arr[end] {
			continue
		}
		arr[i], arr[small] = arr[small], arr[i]
		small++
	}
	arr[small], arr[end] = arr[end], arr[small]
	return small
}
